import bpy
from bpy.types import Operator, Panel
from bpy.props import *

class MATERIAL_OT_CreateNew(bpy.types.Operator):
    """新建材质（可设定名称和颜色）"""
    bl_idname = "material.create_new"
    bl_label = "新建材质"
    bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}

    name: bpy.props.StringProperty(
        name="材质名称",
        default="新材质",
        description="新材质的名称"
    )
    color: bpy.props.FloatVectorProperty(
        name="基础颜色",
        subtype='COLOR',
        size=4,
        default=(1.0, 1.0, 1.0, 1.0),
        min=0.0,
        max=1.0,
        description="材质的基础颜色（RGBA）"
    )

    def invoke(self, context, event):
        return context.window_manager.invoke_props_dialog(self)

    def execute(self, context):
        # 使用用户输入的名称创建材质
        new_mat = bpy.data.materials.new(name=self.name)
        # 设置材质基础颜色（适用于BSDF节点）
        if new_mat.use_nodes:
            bsdf = new_mat.node_tree.nodes.get("Principled BSDF")
            if bsdf:
                bsdf.inputs["Base Color"].default_value = self.color
        else:
            new_mat.diffuse_color = self.color  # 兼容非节点模式
        return {'FINISHED'}

class MATERIAL_OT_ChangeBaseColor(bpy.types.Operator):
    """更改材质基础颜色"""
    bl_idname = "material.change_base_color"
    bl_label = "更改基础色"
    bl_options = {'REGISTER', 'UNDO'}
    color: bpy.props.FloatVectorProperty(
        name="基础颜色",
        subtype='COLOR',
        size=4,
        default=(1.0, 1.0, 1.0, 1.0),
        min=0.0,
        max=1.0,
        description="材质的基础颜色（RGBA）"
    )
    
    def execute(self, context):
        if not (context.active_object and context.active_object.active_material):
            return {'CANCELLED'}
            
        mat = context.active_object.active_material
        
        if not mat.use_nodes:
            mat.diffuse_color = self.color
            return {'FINISHED'}
            
        # 查找所有BSDF类型节点
        bsdf_nodes = [node for node in mat.node_tree.nodes 
                     if node.type.startswith('BSDF')]
        
        if not bsdf_nodes:
            self.report({'ERROR'}, "未找到BSDF节点")
            return {'CANCELLED'}
            
        # 查找颜色输入插槽
        color_socket = None
        for bsdf in bsdf_nodes:
            for input in bsdf.inputs:
                if input.type == 'RGBA':
                    color_socket = input
                    break
            if color_socket:
                break
                
        if not color_socket:
            self.report({'ERROR'}, "未找到颜色输入插槽")
            return {'CANCELLED'}
            
        # 设置颜色值并确保在有效范围内
        color_socket.default_value = (
            max(0.0, min(1.0, self.color[0])),
            max(0.0, min(1.0, self.color[1])),
            max(0.0, min(1.0, self.color[2])),
            1.0
        )
        
        return {'FINISHED'}

class MaterialListPanel(bpy.types.Panel):
    """侧边栏材质列表面板"""
    bl_label = "材质列表"
    bl_idname = "VIEW3D_PT_material_list"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = "Material"

    def draw(self, context):
        layout = self.layout
        # 新建材质按钮
        row = layout.row()
        row.operator("material.create_new", text="新建材质", icon="ADD")
        # 列出所有材质
        for mat in bpy.data.materials:
            if mat.name == "Dots Stroke":
                continue
            row = layout.row()
            icon = mat.preview.icon_id if mat.preview else bpy.types.UILayout.icon(mat)
            op = row.operator("material.drag_assign", text=mat.name, icon_value=icon)
            op.material_name = mat.name

class MATERIAL_OT_DragAssign(bpy.types.Operator):
    """点击按钮分配材质到选中物体"""
    bl_idname = "material.drag_assign"
    bl_label = "分配材质"
    bl_options = {'REGISTER', 'UNDO'}

    material_name: bpy.props.StringProperty()

    @classmethod
    def poll(cls, context):
        return context.object is not None and context.object.type == 'MESH'

    def execute(self, context):
        obj = context.object
        if not obj or not self.material_name:
            return {'CANCELLED'}

        mat = bpy.data.materials.get(self.material_name)
        if not mat:
            self.report({'ERROR'}, f"材质 '{self.material_name}' 不存在")
            return {'CANCELLED'}

        # 如果没有材质槽，则添加一个
        if not obj.material_slots:
            obj.data.materials.append(mat)
        else:
            obj.material_slots[0].material = mat  # 替换第一个材质槽

        return {'FINISHED'}

export_classes = [
    MATERIAL_OT_CreateNew,
    MATERIAL_OT_ChangeBaseColor,
    MATERIAL_OT_DragAssign,
    MaterialListPanel
]